Skip to content

feat(pipeline): Add inference and lineage step types - #6224

Open
Rishabh0255 wants to merge 1 commit into
aws:masterfrom
Rishabh0255:feat/inference-lineage-steps
Open

feat(pipeline): Add inference and lineage step types#6224
Rishabh0255 wants to merge 1 commit into
aws:masterfrom
Rishabh0255:feat/inference-lineage-steps

Conversation

@Rishabh0255

@Rishabh0255 Rishabh0255 commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Add 4 pipeline step classes:

EndpointConfigStep, EndpointStep (SageMaker inference deployment)
InferenceComponentStep (multi-model endpoint support)
LineageStep (ML governance tracking)

Design: each step takes step_args captured under a PipelineSession, following the convention used by TrainingStep and ModelStep. Session.endpoint_from_production_variants, Session.create_endpoint and Session.create_inference_component route through _intercept_create_request -- a pass-through under a plain Session, capturing the request without a service call under a PipelineSession. LineageStep records one entity per step, with step_args from Action/Artifact/Context/Association.create(); an association can reference an entity created by another step via Steps..ActionArns['']. Each step validates step_args provenance with validate_step_args_input, rejecting a wrong producer or a raw dict.

Chaining note: the resource created by an EndpointConfig step carries an execution-unique name suffix, so a downstream step must reference it as config_step.properties.EndpointConfigName rather than the requested name. The Endpoint step does not suffix.

Retryability: only EndpointConfigStep is retryable. Cacheability: EndpointConfigStep and EndpointStep are structurally cacheable via cache_config.

Includes 29 unit tests and two integration tests covering all four step types end to end against the service (2 passed in 473.64s).

Supersedes #6140 -- same feature, reworked to step_args per review feedback.

Issue #, if available:

Description of changes:

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

return _SHAPE_CACHE[cache_key]


def validate_step_arguments(

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

why do we need all of this additional validation here? Other steps do not have this explicit validation. How are these steps different from other steps?

def __init__(
self,
name: str,
arguments: Dict[str, Any],

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not the right implementation for any of these steps. It will be very difficult to construct these arguments manually. We need to use the existing pysdk constructs and pass them as arguments. Please see how Training/Model steps are implemented and follow that pattern here. You must use step_args from PipelineSession instead of raw arguments: dict

SDK primitives exists for all four steps, and it eliminates the entire _argument_validation.py machinery

return get_execution_role()


def test_lineage_step_execute_end_to_end(sagemaker_session, pipeline_session, role):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add integ tests for other steps as well

Adds four pipeline step classes -- EndpointConfigStep, EndpointStep,
InferenceComponentStep and LineageStep -- along with their StepTypeEnum
values and package exports.

Each step takes step_args captured under a PipelineSession, following
the convention used by TrainingStep and ModelStep:

- Session.endpoint_from_production_variants, Session.create_endpoint
  and Session.create_inference_component route their service calls
  through _intercept_create_request. Under a plain Session this is a
  pass-through; under a PipelineSession the request is captured and no
  service call is made.
- LineageStep records one lineage entity per step, with step_args from
  Action.create(), Artifact.create(), Context.create() or
  Association.create(). Record._invoke_api captures those four calls
  under a PipelineSession. An association can reference an entity
  created by another step via Steps.<name>.ActionArns['<entity>'].
- Each step validates the provenance of its step_args with
  validate_step_args_input, rejecting a wrong producer or a raw dict.

Note for chaining: the resource created by an EndpointConfig step
carries an execution-unique name suffix, so a downstream step must
reference it as config_step.properties.EndpointConfigName rather than
the requested name.

Tested with 29 unit tests plus two integration tests that exercise all
four step types end to end against the service.
---
X-AI-Prompt: Add the new inference and lineage pipeline step types to
the public SDK, using step_args captured via PipelineSession per review
feedback, and verify them with integration tests
X-AI-Tool: kiro-cli
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants